Support rescoring with expand_nested_docs on the Lucene engine (#3125) - #3579

TopDocs.merge(k, perLeafTopDocs) may throw when any leaf returned
NestedKnnUtil.EMPTY_TOP_DOCS because merge requires each input to be sorted or
non-empty with consistent structure, and empty leaves also have shardIndex unset.
Ensure every per-leaf TopDocs is non-null and that empty leaves are either filtered
out or produce an empty but valid TopDocs so the merge is safe.

src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [187-197]

 TopDocs[] perLeafTopDocs = indexSearcher.getTaskExecutor().invokeAll(rescoreTasks).toArray(TopDocs[]::new);
+for (int i = 0; i < perLeafTopDocs.length; i++) {
+    if (perLeafTopDocs[i] == null) {
+        perLeafTopDocs[i] = NestedKnnUtil.EMPTY_TOP_DOCS;
+    }
+}
 TopDocs topKParents = TopDocs.merge(k, perLeafTopDocs);
 
 List<Map<Integer, Float>> reducedResults = new ArrayList<>(leafReaderContexts.size());
 for (int i = 0; i < leafReaderContexts.size(); i++) {
     reducedResults.add(new HashMap<>());
 }
 for (ScoreDoc scoreDoc : topKParents.scoreDocs) {
     reducedResults.get(scoreDoc.shardIndex).put(scoreDoc.doc, scoreDoc.score);
 }
 return reducedResults;

The cache key leafReaderContext.ord is only unique within a single IndexReader; if
createWeight is ever called across different readers in the same instance it could
collide. More importantly, the first get is racy with putIfAbsent, which is fine,
but the method may still build the bits twice unnecessarily. Consider
computeIfAbsent with a wrapper to handle the null sentinel, or document that the map
is scoped to one createWeight invocation.

src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [277-278]

+Bits existing = filterBitsByLeaf.putIfAbsent(leafReaderContext.ord, bits);
+return existing != null ? existing : bits;
 
-

When a leaf's leafResult is empty the task returns NestedKnnUtil.EMPTY_TOP_DOCS
whose scoreDocs is a shared empty array, but later code sets scoreDoc.shardIndex =
leafOrd on returned docs. While the empty array has no elements and is safe here,
TopDocs.merge requires each entry to be non-null; verify EMPTY_TOP_DOCS is non-null
(and consider returning a fresh empty TopDocs to avoid any shared-instance mutation
risk if callers ever iterate differently).

src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [153-163]

 rescoreTasks.add(() -> {
     if (leafResult.isEmpty()) {
-        return NestedKnnUtil.EMPTY_TOP_DOCS;
+        return new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]);
     }
     Bits queryFilter = filterBits(filterBitsByLeaf, leafReaderContext, filterWeight);
     DocIdSetIterator allSiblings = queryUtils.getAllSiblings(
         leafReaderContext,
         leafResult.keySet(),
         internalNestedKnnVectorQuery.getParentFilter(),
         queryFilter
     );

NestedKnnUtil.EMPTY_TOP_DOCS is likely a shared singleton; returning it from
multiple leaf tasks and then passing them to TopDocs.merge may be fine, but using a
shared mutable TopDocs (with a shared scoreDocs array) risks subtle bugs if any
downstream code mutates shardIndex or doc. Consider returning a fresh empty TopDocs
per leaf to match the pattern used elsewhere and avoid potential shared-state issues
(similar to the PerLeafResult.EMPTY_RESULT NPE fix noted in CHANGELOG).

src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [154-156]

 rescoreTasks.add(() -> {
     if (leafResult.isEmpty()) {
-        return NestedKnnUtil.EMPTY_TOP_DOCS;
+        return new TopDocs(new org.apache.lucene.search.TotalHits(0, org.apache.lucene.search.TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]);
     }

ConcurrentHashMap does not permit null values, and filterBits already handles null
by early-returning without caching. However, if queryUtils.createBits returns a
non-null Bits whose ord collides across concurrent tasks, putIfAbsent is correct but
the bits computed by the loser are discarded unused. This is fine functionally, but
note that leaf ord is unique per leaf, so collisions only happen when the same leaf
runs rescore and expansion concurrently — ensure the two phases do not run
concurrently (currently they are sequential: rescore completes before retrieveAll),
which is the case, so this is safe.

src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [91]

+Map<Integer, Bits> filterBitsByLeaf = new ConcurrentHashMap<>();
 
-

When a leaf has no candidates, returning NestedKnnUtil.EMPTY_TOP_DOCS (a shared
static instance) is fine for merging, but any downstream code that mutates
scoreDoc.shardIndex on these results would corrupt the shared empty array. Although
scoreDocs is empty here so no mutation occurs, confirm that EMPTY_TOP_DOCS is never
mutated elsewhere; otherwise return a fresh empty TopDocs to be safe.

src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [154-156]

 rescoreTasks.add(() -> {
     if (leafResult.isEmpty()) {
-        return NestedKnnUtil.EMPTY_TOP_DOCS;
+        return new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]);
     }

leafReaderContext.ord is the ordinal within the parent reader. If the query is
executed against sub-readers at different levels (e.g., per-leaf vs. top-level), two
different leaves could share the same ord, causing bit-set collisions. Consider
keying the cache by leafReaderContext.id() or the LeafReaderContext itself to
guarantee uniqueness across all contexts seen by this query.

src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [277-278]

+Bits existing = filterBitsByLeaf.putIfAbsent(leafReaderContext.ord, bits);
+return existing != null ? existing : bits;
 
-

The ScoreDoc instances returned here are mutated with a shardIndex tag and later
used to look up the leaf via scoreDoc.shardIndex. However, TopDocs.merge(k,
perLeafTopDocs) with any non-null shardIndex semantics may rewrite shardIndex on the
returned ScoreDocs (Lucene's merge sets shardIndex when the input docs have
shardIndex set and i is passed). Verify that the shardIndex tag survives
TopDocs.merge — if Lucene overwrites it with the array index, the subsequent
reducedResults.get(scoreDoc.shardIndex) will populate the wrong leaf map and
silently corrupt results. Consider using a side map keyed by identity or encoding
the leaf ord into the doc id.

src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [180-183]

+for (ScoreDoc scoreDoc : leafTopDocs.scoreDocs) {
+    scoreDoc.shardIndex = leafOrd;
+}
+return leafTopDocs;
 
-

resolveExactSearcher() is called from both rescoreToTopKParents and scoreAllSiblings
in createWeight, and each call constructs a new ExactSearcher when exactSearcher is
unset. Cache the resolved instance (e.g., in a local variable passed through, or
memoize once at the start of createWeight) to avoid building two ExactSearcher
instances per query execution, each reaching for the ModelDao singleton.

src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [114-116]

 private ExactSearcher resolveExactSearcher() {
-    return exactSearcher != null ? exactSearcher : new ExactSearcher(ModelDao.OpenSearchKNNModelDao.getInstance());
+    if (exactSearcher != null) {
+        return exactSearcher;
+    }
+    return new ExactSearcher(ModelDao.OpenSearchKNNModelDao.getInstance());
 }